Here are some key observations to aid the review process:

void IrContainer::removeStatementsOwnedBy(const Fusion* fusion) {
  auto vals_it = per_fusion_vals_.find(fusion);
  if (vals_it != per_fusion_vals_.end()) {
    const auto& owned = vals_it->second;
    std::erase_if(vals_up_, [&](const std::unique_ptr<Val>& v) {
      if (owned.count(v.get()) > 0) {
        vals_.erase(v.get());
        return true;
      }
      return false;
    });
    per_fusion_vals_.erase(vals_it);
  }

  auto exprs_it = per_fusion_exprs_.find(fusion);
  if (exprs_it != per_fusion_exprs_.end()) {
    const auto& owned = exprs_it->second;
    std::erase_if(exprs_up_, [&](const std::unique_ptr<Expr>& e) {
      if (owned.count(e.get()) > 0) {
        exprs_.erase(e.get());
        return true;
      }
      return false;
    });
    per_fusion_exprs_.erase(exprs_it);
  }
}

a.ir_container()->transferStatementOwnership(&b, &a);
b.ir_container()->transferStatementOwnership(&a, &b);

if (a.ir_container_) {
  for (auto val : a.vals()) {
    val->ir_container_ = &a;
  }
  for (auto expr : a.deterministic_exprs()) {
    expr->ir_container_ = &a;

int64_t numValsExcludingShortcuts() const noexcept {
  int64_t count = std::ssize(ir_container()->valsOwnedBy(this));
  count -= (zero_val_ != nullptr) + (one_val_ != nullptr) +
      (true_val_ != nullptr) + (false_val_ != nullptr) +
      (magic_zero_val_ != nullptr);
  return count;
}

csrc/fusion.cpp
Shared-container safety gap in rollback

In removeStatementsCreatedAfter, statements are popped from the back of the container-global exprs_up_ deque (line 445), but only the current Fusion's per-fusion map is cleaned up (line 449: c->per_fusion_exprs_[this].erase(e)).

In a shared container scenario where another Fusion has appended statements after the StatementGuard snapshot, the statement at the back of exprs_up_ could belong to the other Fusion. This code would:

1. Destroy the unique_ptr (freeing the Expr memory)
2. Remove it from exprs_ (global set)
3. Only erase from per_fusion_exprs_[this] — leaving a dangling pointer in the other Fusion's per-fusion set

The same issue applies to the vals_up_ loop below (line 467).

This is safe for single-Fusion containers today, but since this PR's stated goal is enabling shared-container safety, the rollback path should also clean up the correct owner's per-fusion map — e.g., by looking up which Fusion actually owns each statement being removed.

csrc/fusion.h
numExprs/numVals still return container-global counts

numExprs() and numVals() forward directly to the container's total counts, not per-Fusion counts. This is inconsistent with the fact that vals(), unordered_exprs(), and deterministic_* now return per-Fusion filtered results.

More importantly, StatementGuard uses numExprs() and numVals() to snapshot statement counts and later passes them to removeStatementsCreatedAfter, which pops from the global deques. In a shared container where another Fusion adds statements concurrently, the snapshot will include the other Fusion's additions, and the rollback will remove them — destroying statements that don't belong to this Fusion.

For single-Fusion containers this is fine, but these methods should be updated to return per-Fusion counts (or at least documented as container-global) before shared containers are actually used.

csrc/ir/container.cpp, line 220
Iterator invalidation UB in transferStatementOwnership

The per_fusion_vals_.find(from) iterator vals_it is stored, and then per_fusion_vals_[to] is called via operator[]. If to is a new key (not yet in the map), the insertion can trigger a rehash, which invalidates all existing iterators — including vals_it. The subsequent to_vals.insert(vals_it->second.begin(), vals_it->second.end()) and per_fusion_vals_.erase(vals_it) both dereference the now-dangling iterator, producing undefined behaviour.

This is actively triggered by Fusion::swap, where the post-swap container has only one key (e.g. &b) and calling operator[&a] inserts a second entry, doubling the bucket count and rehashing.

The same pattern repeats for per_fusion_exprs_ (lines 215–220).

The fix is to extract the source set before any mutation of the map:

auto vals_it = per_fusion_vals_.find(from);
if (vals_it != per_fusion_vals_.end()) {
  // Copy the source set before operator[] can rehash and invalidate vals_it.
  auto from_vals = std::move(vals_it->second);
  per_fusion_vals_.erase(vals_it);
  per_fusion_vals_[to].insert(from_vals.begin(), from_vals.end());
}

auto exprs_it = per_fusion_exprs_.find(from);
if (exprs_it != per_fusion_exprs_.end()) {
  auto from_exprs = std::move(exprs_it->second);
  per_fusion_exprs_.erase(exprs_it);
  per_fusion_exprs_[to].insert(from_exprs.begin(), from_exprs.end());
}

csrc/fusion.cpp, line 444
LIFO invariant breaks silently under shared containers

The expr-removal loop unconditionally pops from c->exprs_up_.back() and asserts (hard-abort via NVF_ERROR) that the tail entry belongs to this Fusion. The same invariant governs the vals loop (lines 447–466).

For a single-Fusion container this holds, but once two Fusions share a container (the stated goal of Phase 2), another Fusion may register statements into the shared deque after this StatementGuard scope opened. In that case the tail could belong to the other Fusion, NVF_ERROR fires, and the process aborts rather than correctly rolling back only this Fusion's statements.

Effectively, StatementGuard becomes incompatible with concurrent statement creation across Fusions sharing a container. Since this is Phase 2 infrastructure that explicitly targets shared containers, this is worth surfacing early — even if addressed in a later PR in the chain.

One possible approach for future-proofing: instead of a deque tail pop, iterate c->exprs_up_ in reverse and erase only entries owned by this, or collect the set of exprs created during the scope and remove them directly by pointer rather than relying on positional ordering.

csrc/ir/container.cpp, line 215
Iterator invalidated by insertion before use

per_fusion_vals_[to] performs an insertion when to is not yet a key in the map. Per the C++ standard, any insertion that triggers a rehash invalidates all existing iterators. Since vals_it was obtained before this call, if the insertion causes a rehash, both the subsequent vals_it->second read and the per_fusion_vals_.erase(vals_it) call are undefined behaviour.

The same hazard applies in the per_fusion_exprs_ block a few lines below, where exprs_it is similarly invalidated by per_fusion_exprs_[to].

A safe alternative is to extract the node first, which is both iterator-safe and avoids copying the entire set:

void IrContainer::transferStatementOwnership(
    const Fusion* from,
    const Fusion* to) {
  auto vals_node = per_fusion_vals_.extract(from);
  if (vals_node) {
    auto& to_vals = per_fusion_vals_[to];
    to_vals.insert(
        std::make_move_iterator(vals_node.mapped().begin()),
        std::make_move_iterator(vals_node.mapped().end()));
  }

  auto exprs_node = per_fusion_exprs_.extract(from);
  if (exprs_node) {
    auto& to_exprs = per_fusion_exprs_[to];
    to_exprs.insert(
        std::make_move_iterator(exprs_node.mapped().begin()),
        std::make_move_iterator(exprs_node.mapped().end()));
  }
}

std::unordered_map::extract removes the node without rehashing the table, so no other iterators are invalidated, and the existing set memory is moved rather than copied.

csrc/ir/container.cpp, line 244
Missing Val::uses_ cleanup when removing exprs

removeStatementsCreatedAfter (the StatementGuard rollback path) carefully calls in->removeUse(e) for every input Val before erasing each Expr. removeStatementsOwnedBy (the Fusion::clear() path) does not perform this cleanup.

For single-Fusion containers this is currently safe because every owned Val is also destroyed in the same call, so dangling uses_ entries are never read. However, for shared containers — the stated end-goal of the Phase 2 chain — it is entirely possible that a Val owned by Fusion B references (via its uses_ set) an Expr owned by Fusion A. When Fusion A is cleared, Fusion B's Val retains a now-dangling pointer in its uses_ set, making any subsequent traversal of Val::uses_ unsafe.

Removing exprs first (before destroying vals) and calling in->removeUse(e) for each input would make this safe for the multi-Fusion case and keep the two removal paths consistent with each other.

csrc/fusion.cpp, line 436
operator[] in assertion silently mutates the map

c->per_fusion_exprs_[this] calls the non-const operator[], which inserts a default-constructed empty set if this is not yet a key in per_fusion_exprs_. Logically this should never fire (we only enter the loop body when exprsOwnedBy(this) is non-empty, so this must already be present), but the operator[] call still mutates the map as a side-effect of running the assertion.

The same pattern appears in the val loop on line 448 (c->per_fusion_vals_[this].count(v)).

Prefer a read-only lookup to avoid the accidental mutation:

    NVF_ERROR(
        c->per_fusion_exprs_.count(this) > 0 &&
            c->per_fusion_exprs_.at(this).count(e) > 0,
        "removeStatementsCreatedAfter: tail expr belongs to another Fusion");

csrc/ir/container.cpp, line 216
Iterator invalidation UB in transferStatementOwnership

per_fusion_vals_[to] (line 205) uses operator[], which inserts a new entry when to is not yet in the map. Inserting into an std::unordered_map may trigger a rehash, which invalidates all iterators — including vals_it obtained on line 203. The subsequent vals_it->second.begin() / .end() call on line 206 is therefore undefined behaviour whenever to was not previously in the map. The same issue affects exprs_it / per_fusion_exprs_[to] on lines 212–214.

Safe fix — extract the from set before touching the map for to:

void IrContainer::transferStatementOwnership(
    const Fusion* from,
    const Fusion* to) {
  auto vals_it = per_fusion_vals_.find(from);
  if (vals_it != per_fusion_vals_.end()) {
    // Move out before any insertion that could rehash
    auto from_set = std::move(vals_it->second);
    per_fusion_vals_.erase(vals_it);
    per_fusion_vals_[to].insert(from_set.begin(), from_set.end());
  }

  auto exprs_it = per_fusion_exprs_.find(from);
  if (exprs_it != per_fusion_exprs_.end()) {
    auto from_set = std::move(exprs_it->second);
    per_fusion_exprs_.erase(exprs_it);
    per_fusion_exprs_[to].insert(from_set.begin(), from_set.end());
  }
}

Alternatively, std::unordered_map::extract can be used — it transfers the node without rehashing, keeping all iterators valid.

csrc/ir/container.cpp, line 244
Missing Val::uses_ cleanup when removing owned expressions

removeStatementsOwnedBy destroys exprs via std::erase_if (which calls unique_ptr destructors) but never removes those exprs from the uses_ lists of their input Vals. Compare with removeStatementsCreatedAfter, which explicitly calls in->removeUse(e) for every input before popping each expr.

For a shared container this is a latent bug: if a Val owned by Fusion B has Fusion A's Expr in its uses_ list, and Fusion::clear() for Fusion A calls removeStatementsOwnedBy, Fusion B's Val is left with a dangling pointer in uses_.

For the current single-Fusion use case all Vals are destroyed in the same call, so the stale uses_ pointers are never read. However, given that the stated goal of Phase 2 is shared containers, the cleanup should mirror removeStatementsCreatedAfter:

std::erase_if(exprs_up_, [&](const std::unique_ptr<Expr>& e) {
  if (owned.count(e.get()) > 0) {
    for (Val* in : e->inputs()) {
      in->removeUse(e.get());
    }
    exprs_.erase(e.get());
    return true;
  }
  return false;
});

csrc/fusion.cpp, line 466
LIFO invariant holds only for single-Fusion containers

Both loops pop from the global exprs_up_ / vals_up_ deque tail and then assert the popped element belongs to this Fusion. This is sound for a single-Fusion container, but the invariant breaks as soon as two Fusions interleave statement registration in a shared container.

Consider:

Guard created for Fusion A  (num_exprs_before = N)
Fusion A registers expr_a   → tail of exprs_up_ is expr_a
Fusion B registers expr_b   → tail of exprs_up_ is expr_b   ← different Fusion!
Guard for A is destroyed:
  exprsOwnedBy(A).size() > N  → true
  exprs_up_.back() == expr_b  → NVF_ERROR fires

The NVF_ERROR will surface the violation rather than silently corrupt state, which is good. But since this invariant is a load-bearing assumption for the StatementGuard rollback mechanism, it is worth recording it explicitly (e.g., a NVF_CHECK(sharingCount() == 1, ...) guard at the top of the function) or documenting that StatementGuard usage must be restricted to single-Fusion containers until a per-Fusion deque is introduced.

csrc/ir/container.cpp, line 243
removeStatementsOwnedBy silently skips Val::uses_ cleanup

When removing expressions via removeStatementsOwnedBy, the inputs' uses_ sets are never updated. In contrast, removeStatementsCreatedAfter (the StatementGuard path) explicitly calls in->removeUse(e) for each input:

// removeStatementsCreatedAfter (correct):
for (Val* in : e->inputs()) {
  in->removeUse(e);
}

For the current single-Fusion case (Fusion::clear()), this is safe because all vals owned by this Fusion are destroyed in the same call — no live Val* retains the dangling uses_ entry. However, once shared containers land in a later phase, a Val owned by Fusion B could be an input to an Expr owned by Fusion A. After removeStatementsOwnedBy(A) destroys that expr, Fusion B's val retains a dangling pointer in its uses_ set, which any subsequent val->uses() traversal will dereference.

Consider adding the same cleanup here to make the invariant hold regardless of sharing:

std::erase_if(exprs_up_, [&](const std::unique_ptr<Expr>& e) {
  if (owned.count(e.get()) > 0) {
    for (Val* in : e->inputs()) {
      in->removeUse(e.get());
    }
    exprs_.erase(e.get());
    return true;
  }
  return false;
});

csrc/fusion.cpp, line 443
LIFO invariant breaks silently under shared containers with interleaved allocation

removeStatementsCreatedAfter pops from the back of the global exprs_up_ deque and then asserts the popped element is owned by this:

Expr* e = c->exprs_up_.back().get();
NVF_ERROR(
    c->per_fusion_exprs_[this].count(e) > 0,
    "removeStatementsCreatedAfter: tail expr belongs to another Fusion");

This is safe today (single-Fusion containers) but fails in the target Phase 2 multi-Fusion scenario. Consider Fusion A and B sharing a container where StatementGuard is nested or interleaved:

exprs_up_ (global tail): [..., A1, B1, A2]

When A's guard destructs, it checks exprs_up_.back() → A2 ✓. But if B's guard destructs first, it sees A2 at the back → NVF_ERROR fires. There is no safe ordering guarantee once allocations interleave.

The NVF_ERROR prevents silent corruption, which is good, but it means StatementGuard cannot be used safely from two Fusions sharing a container. The same issue applies to the vals_up_ loop. A comment calling this out explicitly (or a NVF_ERROR early in the function checking !ir_container()->hasMultipleFusions()) would help future readers understand the precondition before relaxing it.

csrc/fusion.h, line 570
numValsExcludingShortcuts misses typed shortcut vals (zeroVal(dtype), oneVal(dtype))

The function subtracts only the 5 named shortcut fields (zero_val_, one_val_, true_val_, false_val_, magic_zero_val_). Typed shortcuts created via zeroVal(DataType) and oneVal(DataType) are stored separately (e.g., in managed_data_) and are also registered in per_fusion_vals_ — but they are not excluded from this count.

If a typed shortcut val (e.g., zeroVal(DataType::Float)) is lazily created inside a StatementGuard scope, it will be appended to the back of vals_up_ and included in numValsExcludingShortcuts(). The removeStatementsCreatedAfter loop will then pop and destroy it, leaving managed_data_ holding a dangling Val* for the rest of the Fusion's lifetime. Any subsequent use of that typed constant (e.g., during lowering) would dereference freed memory.

The existing axioms() call in the StatementGuard constructor is the pattern used to pre-initialize lazy state before the guard captures the baseline. Consider extending that pattern to also pre-initialize any typed shortcut vals that callers are likely to create inside the guard, or document the precondition that typed shortcut vals must not be first-created within a StatementGuard scope.